import json
import yaml

from playwright.sync_api import APIRequestContext
from assertpy import assert_that

from .constants import NODE_PATH, COLOCATION_CONFIG_CONFIGMAP, VOLCANO_CONFIG_CONFIGMAP


def fetch_all_node(api_context: APIRequestContext) -> dict:
    response = api_context.get(NODE_PATH)
    assert_that(response.ok, f'Failed to fetch nodes: {response.text()}')
    return response.json()
    

def fetch_node(api_context: APIRequestContext, node_name: str) -> dict:
    response = api_context.get(f'{NODE_PATH}{node_name}')
    assert_that(response.ok, f'Failed to fetch node {node_name}: {response.text()}')
    return response.json()


def fetch_colocation_config(api_context: APIRequestContext) -> tuple[dict, dict, dict]:
    response = api_context.get(COLOCATION_CONFIG_CONFIGMAP)
    assert_that(response.ok, f'Failed to fetch colocation config: {response.text()}')
    config_data = response.json()
    colocation_options = json.loads(config_data['data']['colocation-options'])
    volcano_scheduler_options = json.loads(config_data['data']['volcano-scheduler-options'])
    rubik_options = json.loads(config_data['data']['rubik-options'])
    return colocation_options, volcano_scheduler_options, rubik_options


def fetch_volcano_config(api_context: APIRequestContext) -> dict:
    response = api_context.get(VOLCANO_CONFIG_CONFIGMAP)
    assert_that(response.ok, f'Failed to fetch volcano config: {response.text()}')
    config_data = response.json()
    volcano_scheduler_config = yaml.safe_load(config_data['data']['volcano-scheduler.conf'])
    return volcano_scheduler_config
